import pandas as pd
import datetime
import baostock as bs
import sys
bs.login()

day1 = datetime.timedelta(days = 1)

def fmt_dt(dt):
    return dt.strftime("%Y-%m-%d")


code_df = pd.read_csv("./data/hs300.csv")
df = pd.read_csv("./data/history_k/merge/hs300/hs300.csv")
code_map = {}
df_map = {}
for index, row in code_df.iterrows():
    code_map[row.code] = row.code_name
    df_map[row.code] = df[df.code == row.code]



codes = list(code_map.keys())



def is_dfx(code, dt):
    cdf = df_map[code]

    d3 = dt
    d3_index = tdays.index(d3)
    d2 = tdays[d3_index - 1]
    d1 = tdays[d3_index - 2]
    # print(d1, d2, d3)

    r1 = cdf[cdf.date == d1]
    r2 = cdf[cdf.date == d2]
    r3 = cdf[cdf.date == d3]

    l1 = r1.low.values[0]
    l2 = r2.low.values[0]
    l3 = r3.low.values[0]


    h1 = r1.high.values[0]
    h2 = r2.high.values[0]
    h3 = r3.high.values[0]
    return h2 < h1 and h2 < h3 and l2 < l1 and l2 < l3



cdf = df_map[codes[0]]
td = datetime.datetime.today()


rs = bs.query_trade_dates(start_date=fmt_dt(td - day1 * 300), end_date=fmt_dt(td))
tdays = []
for index, row in rs.get_data().iterrows():
    if row['is_trading_day'] == '1':
        tdays.append(row['calendar_date'])


if len(sys.argv) == 2:
    day =  sys.argv[1]
else:
    day = tdays[-1]
for code in codes:
    #datetime.datetime.strptime("2021-05-14", "%Y-%m-%d")
    flag = is_dfx(code, day)
    #print(code, code_map[code])
    if flag:
        print(code, code_map[code])


